package com.xiaobias;

import com.rabbitmq.client.*;

import java.nio.charset.StandardCharsets;
import java.sql.*;
import java.util.Set;
import java.util.concurrent.CopyOnWriteArraySet;

public class PasswordConsumer {

    private static final String QUEUE_NAME = "password_queue";
    private static final String MQ_HOST = "localhost";
    private static final int MQ_PORT = 5672;
    private static final String MQ_USERNAME = "admin";
    private static final String MQ_PASSWORD = "admin";

    private static final String DB_URL = "jdbc:mysql://localhost:3306/password_db";
    private static final String DB_USER = "root";
    private static final String DB_PASSWORD = "root123";

    private static final int BATCH_SIZE = 1000; // 每批插入的数量

    public static void main(String[] args) throws Exception {
        // 初始化 RabbitMQ 连接
        ConnectionFactory factory = new ConnectionFactory();
        factory.setHost(MQ_HOST);
        factory.setPort(MQ_PORT);
        factory.setUsername(MQ_USERNAME);
        factory.setPassword(MQ_PASSWORD);

        com.rabbitmq.client.Connection connection = factory.newConnection();
        Channel channel = connection.createChannel();

        // 声明队列
        channel.queueDeclare(QUEUE_NAME, false, false, false, null);
        System.out.println("Waiting for messages. To exit, press Ctrl+C");

        // 初始化 MySQL 连接
        java.sql.Connection dbConnection = DriverManager.getConnection(DB_URL, DB_USER, DB_PASSWORD);
        dbConnection.setAutoCommit(false); // 关闭自动提交，启用批量插入
        dropAndCreateTable(dbConnection);// 创建表

        // 用于缓存密码的列表
        Set<String> passwordBatch = new CopyOnWriteArraySet<>();

        // 定义消息处理逻辑
        DeliverCallback deliverCallback = (consumerTag, delivery) -> {
            try {
                String password = new String(delivery.getBody(), StandardCharsets.UTF_8);
                passwordBatch.add(password);
                // 确认消息
                channel.basicAck(delivery.getEnvelope().getDeliveryTag(), false);
                // 如果达到批量大小，插入到数据库
                if (passwordBatch.size() >= BATCH_SIZE) {
                    saveBatchToDatabase(dbConnection, passwordBatch);
                    passwordBatch.clear(); // 清空缓存
                }
            } catch (Exception e) {
                // 处理异常，拒绝消息并重新入队
                channel.basicNack(delivery.getEnvelope().getDeliveryTag(), false, true);
            }
        };

        // 开始消费消息
        channel.basicConsume(QUEUE_NAME, false, deliverCallback, consumerTag -> {
        });

        // 定期检查并插入剩余的密码
        while (true) {
            if (!passwordBatch.isEmpty()) {
                saveBatchToDatabase(dbConnection, passwordBatch);
                passwordBatch.clear();
            }
            Thread.sleep(1000); // 每秒检查一次
        }

    }

    /**
     * 将一批密码保存到数据库
     */
    private static void saveBatchToDatabase(java.sql.Connection dbConnection, Set<String> passwordBatch) {
        StringBuilder sqlSB = null;
        try {
//            System.out.println(Thread.currentThread().getName());
            synchronized (passwordBatch) {
                if (passwordBatch.isEmpty()) {
                    return;
                }
                // 准备 SQL 语句
                sqlSB = new StringBuilder("INSERT INTO `passwords` (password) VALUES");
                for (String password : passwordBatch) {
                    sqlSB.append("('").append(password).append("')").append(",");
                }
                sqlSB.deleteCharAt(sqlSB.length() - 1);
//            sqlSB.append("ON DUPLICATE KEY UPDATE password = password");
                PreparedStatement preparedStatement = dbConnection.prepareStatement(sqlSB.toString());
//                System.out.println(sqlSB.toString());
                preparedStatement.execute(); // 执行批处理
                dbConnection.commit(); // 提交事务
                System.out.println("Inserted " + passwordBatch.size() + " passwords into the database.");
            }
        } catch (SQLIntegrityConstraintViolationException e) {
            try {
                dbConnection.rollback(); // 回滚事务
            } catch (SQLException ex) {
                ex.printStackTrace();
            }

            // 重试
            String duplicate = e.getMessage().replace("Duplicate entry '", "")
                    .replace("' for key 'passwords.PRIMARY'", "");
            passwordBatch.remove(duplicate);
            saveBatchToDatabase(dbConnection, passwordBatch);

        } catch (SQLException e) {
            System.err.println(sqlSB.toString());
            System.err.println(passwordBatch);
            e.printStackTrace();
            try {
                dbConnection.rollback(); // 回滚事务
            } catch (SQLException ex) {
                ex.printStackTrace();
            }
        }
    }

    public static void dropAndCreateTable(java.sql.Connection conn) throws SQLException {
        try (Statement stmt = conn.createStatement()) {
            // 删除表（如果存在）的 SQL 语句
            String dropTableSQL = "DROP TABLE IF EXISTS passwords";
            stmt.executeUpdate(dropTableSQL);
            System.out.println("Table dropped successfully.");

            // 创建新表的 SQL 语句
            String createTableSQL = "CREATE TABLE passwords\n" +
                    "(\n" +
                    "    password VARCHAR(255) primary key\n" +
                    ") collate utf8mb4_bin";
            stmt.executeUpdate(createTableSQL);
            System.out.println("Table created successfully.");
        } catch (SQLException e) {
            System.out.println("Error in drop and create table: " + e.getMessage());
            throw e; // 可以选择重新抛出异常或根据需要处理异常
        }
    }
}